import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, random_split
from utils import LeapGestRecogData
from utils.models import LeapGestureClassifier
from utils.transform import transform
from sklearn.metrics import confusion_matrix, classification_report, roc_curve, auc
import matplotlib.pyplot as plt

def app_main():
    # 创建数据加载器
    dataset = LeapGestRecogData("data/leapGestRecog", transform)
    train_size = int(len(dataset) * 0.8)
    splits = (train_size, len(dataset) - train_size)
    trainset, testset = random_split(dataset, splits)

    trainloader = DataLoader(trainset, batch_size=16, shuffle=True)
    testloader = DataLoader(testset, batch_size=16, shuffle=False)

    # 实例化模型
    model = LeapGestureClassifier()

    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=1e-8)

    # 确保模型保存目录存在
    os.makedirs('models', exist_ok=True)

    # 训练模型
    num_epochs = 20
    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0
        for i, (inputs, labels) in enumerate(trainloader):
            outputs = model(inputs)
            loss = criterion(outputs, labels)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            running_loss += loss.item()

            if (i + 1) % 10 == 0:
                print(f'Epoch [{epoch + 1}/{num_epochs}], Step [{i + 1}/{len(trainloader)}], Loss: {running_loss / (i + 1):.4f}')

    print('Finished Training')

    # 测试模型
    model.eval()#将模型设置为评估模式
    all_labels = []  # 真实标签
    all_predictions = []  # 预测标签
    all_probs = []  # 预测概率

    with torch.no_grad():#上下文用于暂时禁用梯度计算，因为在测试/评估阶段我们不需要计算梯度。
        for inputs, labels in testloader:
            outputs = model(inputs)
            probabilities = torch.softmax(outputs, dim=1)  # 将输出转换为概率
            _, predicted = torch.max(outputs, 1)

            all_predictions.extend(predicted.view(-1).tolist())
            all_labels.extend(labels.view(-1).tolist())
            all_probs.extend(probabilities.tolist())  # 将概率列表转换为Python列表

    # 计算混淆矩阵
    conf_matrix = confusion_matrix(all_labels, all_predictions)
    print("Confusion Matrix:")
    print(conf_matrix)

    # 打印分类报告
    class_report = classification_report(all_labels, all_predictions,zero_division=1)
    print("\nClassification Report:")
    print(class_report)


    # 计算每个类别的 ROC 曲线和 AUC
    fpr = dict()
    tpr = dict()
    roc_auc = dict()
    n_classes = len(set(all_labels))  # 假设类别标签从0开始连续编号

    for i in range(n_classes):
        fpr[i], tpr[i], _ = roc_curve(all_labels, [p[i] for p in all_probs], pos_label=i)
        roc_auc[i] = auc(fpr[i], tpr[i])

    # 绘制所有 ROC 曲线
    plt.figure()
    colors = ['aqua', 'b', 'g', 'r', 'c', 'm', 'y', 'k']
    for i, color in zip(range(n_classes), colors):
        plt.plot(fpr[i], tpr[i], color=color, lw=2,
             label='ROC curve of class {0} (auc = {1:0.2f})'
             ''.format(i, roc_auc[i]))

    plt.plot([0, 1], [0, 1], 'k--', lw=2)
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.5])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('Receiver Operating Characteristic')
    plt.legend(loc="lower right")
    plt.show()

    # 保存整个模型
    torch.save(model, 'models/leap_gesture_model_full.pth')
    print("Saved entire model to 'models/ap_gesture_model_full.pth'")

if __name__ == "__main__":
    app_main()